import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from category_encoders import TargetEncoder
from imblearn.over_sampling import SMOTE

def feature_engineering(X_train, X_test, y_train):
    """对训练集和测试集执行特征工程，并返回标准化后的特征矩阵"""

    X_train = X_train.copy()
    X_test = X_test.copy()

    # ========== 1. 构造比例类特征 ==========
    # X_train['YearsAtCompany_Age_ratio'] = X_train['YearsAtCompany'] / (X_train['Age'] + 1e-5)
    # X_test['YearsAtCompany_Age_ratio'] = X_test['YearsAtCompany'] / (X_test['Age'] + 1e-5)

    # X_train['TotalWorkingYears_Age_ratio'] = X_train['TotalWorkingYears'] / (X_train['Age'] + 1e-5)
    # X_test['TotalWorkingYears_Age_ratio'] = X_test['TotalWorkingYears'] / (X_test['Age'] + 1e-5)
    #
    X_train['CurrentRole_Company_ratio'] = X_train['YearsInCurrentRole'] / (X_train['YearsAtCompany'] + 1e-5)
    X_test['CurrentRole_Company_ratio'] = X_test['YearsInCurrentRole'] / (X_test['YearsAtCompany'] + 1e-5)

    # ========== 2. 收入相关 ==========
    X_train['Income_per_Year'] = X_train['MonthlyIncome'] / (X_train['TotalWorkingYears'] + 1)
    X_test['Income_per_Year'] = X_test['MonthlyIncome'] / (X_test['TotalWorkingYears'] + 1)

    X_train['Income_per_Level'] = X_train['MonthlyIncome'] / (X_train['JobLevel'] + 1e-5)
    X_test['Income_per_Level'] = X_test['MonthlyIncome'] / (X_test['JobLevel'] + 1e-5)



    # ======== 岗位时长与当前岗位经理共事时长 ========
    # X_train['Role_Manager_Diff'] = X_train['YearsInCurrentRole'] - X_train['YearsWithCurrManager']
    # X_test['Role_Manager_Diff'] = X_test['YearsInCurrentRole'] - X_test['YearsWithCurrManager']
    # 比值（加1防止除0）
    # X_train['Role_Manager_Ratio'] = X_train['YearsInCurrentRole'] / (X_train['YearsWithCurrManager'] + 1)
    # X_test['Role_Manager_Ratio'] = X_test['YearsInCurrentRole'] / (X_test['YearsWithCurrManager'] + 1)

    # ========== 3. 时间跨度特征 ==========
    # X_train['Promotion_Frequency'] = X_train['YearsSinceLastPromotion'] / (X_train['YearsAtCompany'] + 1)
    # X_test['Promotion_Frequency'] = X_test['YearsSinceLastPromotion'] / (X_test['YearsAtCompany'] + 1)
    #
    # X_train['RoleChange_Frequency'] = X_train['YearsInCurrentRole'] / (X_train['TotalWorkingYears'] + 1)
    # X_test['RoleChange_Frequency'] = X_test['YearsInCurrentRole'] / (X_test['TotalWorkingYears'] + 1)

    # ========== 4. 满意度组合 ==========
    satisfaction_cols = ['EnvironmentSatisfaction', 'JobSatisfaction', 'RelationshipSatisfaction']
    X_train['Satisfaction_Mean'] = X_train[satisfaction_cols].mean(axis=1)
    X_test['Satisfaction_Mean'] = X_test[satisfaction_cols].mean(axis=1)

    # X_train['Satisfaction_Std'] = X_train[satisfaction_cols].std(axis=1)
    # X_test['Satisfaction_Std'] = X_test[satisfaction_cols].std(axis=1)

    import numpy as np
    import pandas as pd
    from sklearn.preprocessing import StandardScaler

    def feature_engineering(X_train, X_test):
        """对训练集和测试集执行特征工程，并返回标准化后的特征矩阵"""

        # --- 比例类特征 ---
        X_train['YearsAtCompany_Age_ratio'] = X_train['YearsAtCompany'] / (X_train['Age'] + 1e-5)
        X_test['YearsAtCompany_Age_ratio'] = X_test['YearsAtCompany'] / (X_test['Age'] + 1e-5)

        X_train['TotalWorkingYears_Age_ratio'] = X_train['TotalWorkingYears'] / (X_train['Age'] + 1e-5)
        X_test['TotalWorkingYears_Age_ratio'] = X_test['TotalWorkingYears'] / (X_test['Age'] + 1e-5)

        X_train['CurrentRole_Company_ratio'] = X_train['YearsInCurrentRole'] / (X_train['YearsAtCompany'] + 1e-5)
        X_test['CurrentRole_Company_ratio'] = X_test['YearsInCurrentRole'] / (X_test['YearsAtCompany'] + 1e-5)

        # --- 收入相关 ---
        X_train['Income_per_Year'] = X_train['MonthlyIncome'] / (X_train['TotalWorkingYears'] + 1)
        X_test['Income_per_Year'] = X_test['MonthlyIncome'] / (X_test['TotalWorkingYears'] + 1)

        X_train['Income_per_Level'] = X_train['MonthlyIncome'] / (X_train['JobLevel'] + 1e-5)
        X_test['Income_per_Level'] = X_test['MonthlyIncome'] / (X_test['JobLevel'] + 1e-5)

        # --- 时间跨度特征 ---
        X_train['Promotion_Frequency'] = X_train['YearsSinceLastPromotion'] / (X_train['YearsAtCompany'] + 1)
        X_test['Promotion_Frequency'] = X_test['YearsSinceLastPromotion'] / (X_test['YearsAtCompany'] + 1)

        X_train['RoleChange_Frequency'] = X_train['YearsInCurrentRole'] / (X_train['TotalWorkingYears'] + 1)
        X_test['RoleChange_Frequency'] = X_test['YearsInCurrentRole'] / (X_test['TotalWorkingYears'] + 1)

        # --- 满意度组合 ---
        satisfaction_cols = ['EnvironmentSatisfaction', 'JobSatisfaction', 'RelationshipSatisfaction']
        X_train['Satisfaction_Mean'] = X_train[satisfaction_cols].mean(axis=1)
        X_test['Satisfaction_Mean'] = X_test[satisfaction_cols].mean(axis=1)

        X_train['Satisfaction_Std'] = X_train[satisfaction_cols].std(axis=1)
        X_test['Satisfaction_Std'] = X_test[satisfaction_cols].std(axis=1)

        # =====================================================
        # ⭐ 重点增加的交互特征（你列的5个）
        # =====================================================

        # 1. 体现员工职业生涯中在本公司的占比
        X_train['Company_Experience_Ratio'] = X_train['YearsAtCompany'] / (X_train['TotalWorkingYears'] + 1)
        X_test['Company_Experience_Ratio'] = X_test['YearsAtCompany'] / (X_test['TotalWorkingYears'] + 1)

        # 2. 晋升频率（越大表示晋升慢）—— Promotion_Frequency 已有，这里保留即可

        # 3. 当前岗位占公司年限比例
        X_train['CurrentRole_CompanyRatio'] = X_train['YearsInCurrentRole'] / (X_train['YearsAtCompany'] + 1)
        X_test['CurrentRole_CompanyRatio'] = X_test['YearsInCurrentRole'] / (X_test['YearsAtCompany'] + 1)

        # 4. 收入与经验的匹配度
        X_train['Income_Experience_Ratio'] = X_train['MonthlyIncome'] / (X_train['TotalWorkingYears'] + 1)
        X_test['Income_Experience_Ratio'] = X_test['MonthlyIncome'] / (X_test['TotalWorkingYears'] + 1)

        # 5. 住得远且加班多 → 高流失风险
        X_train['Distance_Overtime'] = X_train['DistanceFromHome'] * (X_train['OverTime'] == 'Yes').astype(int)
        X_test['Distance_Overtime'] = X_test['DistanceFromHome'] * (X_test['OverTime'] == 'Yes').astype(int)

    # ========== 5. 对数变换减少偏度 ==========
    for col in ['MonthlyIncome', 'DistanceFromHome', 'NumCompaniesWorked']:
        if col in X_train.columns:
            X_train[col] = np.log1p(X_train[col])
            X_test[col] = np.log1p(X_test[col])

    # ========== 6. 分类特征编码（使用 Target Encoding） ==========
    categorical_features = X_train.select_dtypes(include=['object']).columns
    te = TargetEncoder(cols=categorical_features)
    X_train_categorical = te.fit_transform(X_train[categorical_features], y_train)
    X_test_categorical = te.transform(X_test[categorical_features])

    # 数值特征
    numerical_features = X_train.select_dtypes(exclude=['object']).columns
    X_train_numerical = X_train[numerical_features]
    X_test_numerical = X_test[numerical_features]

    # 合并
    X_train_all = pd.concat([X_train_numerical, X_train_categorical], axis=1)
    X_test_all = pd.concat([X_test_numerical, X_test_categorical], axis=1)

    # ========== 7. 类别不平衡处理（SMOTE过采样） ==========
    smote = SMOTE(random_state=42)
    X_train_resampled, y_train_resampled = smote.fit_resample(X_train_all, y_train)

    # ========== 8. 标准化 ==========
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train_resampled)
    X_test_scaled = scaler.transform(X_test_all)

    return X_train_scaled, X_test_scaled, y_train_resampled
